import torch
from torch.utils.data import Dataset
from torchvision.transforms import transforms
import cv2
import numpy as np
import random
import os
from copy import deepcopy

# class ReSize(torch.nn.Module):
#     def __init__(self, size):
#         super().__init__()
#         # 图片输入大小
#         self.size = size

#     def forward(self, img):
#         h, w = img.shape[0:2]
#         long_border = max(h, w)

#         pad_img = np.full((long_border, long_border, 3), fill_value=114, dtype=np.uint8)
#         pad_img[int((long_border - h)/2):h+int((long_border - h)/2), int((long_border - w)/2):w+int((long_border - w)/2), :] = img
#         pad_img = cv2.resize(pad_img, (self.size, self.size))

#         return pad_img



class MyDataset(Dataset):
    def __init__(self, img_rootPath, mode='train', scale=0.9):
        super().__init__()

        # 数据集划分
        data_dir = list(os.listdir(img_rootPath))
        if mode == 'train':
            # data_dir = data_dir[:int(len(data_dir)*0.3)]
            data_dir = data_dir[:int(len(data_dir)*scale)]
        elif mode == 'test':
            # data_dir = data_dir[:int(len(data_dir)*0.3)]
            data_dir = data_dir[int(len(data_dir)*(scale)):]

        self.img_path_arr = []
        for img_name in data_dir:
            self.img_path_arr.append(os.path.join(img_rootPath, img_name))

        self.size = 256
        self.mask_num = 5
        self.max_mask_size = int(self.size/2)
        self.min_mask_size = int(self.size/4)

        self.transform = transforms.Compose([
            transforms.ToTensor()
        ])
    
    def __getitem__(self, item):
        # 读取图片
        img_path = self.img_path_arr[item]
        img = cv2.imread(img_path)
        img = cv2.resize(img, [self.size, self.size])

        inp_img = deepcopy(img)
        mask = np.zeros([inp_img.shape[0], inp_img.shape[1], 1], dtype=np.uint8)
        for _ in range(self.mask_num):
            rand_x = random.randint(0, self.max_mask_size-1)
            rand_y = random.randint(0, self.max_mask_size-1)
            rand_w = random.randint(self.min_mask_size, self.max_mask_size)
            rand_h = random.randint(self.min_mask_size, self.max_mask_size)
            
            inp_img[rand_y:rand_y+rand_h, rand_x:rand_x+rand_w, :] = 0
            mask[rand_y:rand_y+rand_h, rand_x:rand_x+rand_w, :] = 255

        # 转Tensor并归一化
        X_img = self.transform(inp_img)
        mask = self.transform(mask)
        traget_img = self.transform(img)
        # local_area = torch.tensor([
        #     rand_x, 
        #     rand_x+int(self.size/2),
        #     rand_y,
        #     rand_y+int(self.size/2),
        # ], dtype=torch.int)

        return X_img, mask, traget_img

    def __len__(self):
        return len(self.img_path_arr)


if __name__ == '__main__':
    dataset = MyDataset(r'D:\VOCtrainval_11-May-2012\JPEGImages')
    print(dataset[5][0])